import operator
import pickle
from text_classifier.old_class_to_new_class import class_old_class_new_dict
# DEFAULT_CLASS = 'overall'


import operator
import pickle
from sklearn.svm import LinearSVC
from sklearn.feature_extraction.text import CountVectorizer
import numpy as np
# from text_classifier.old_class_to_new_class import class_old_class_new_dict
# DEFAULT_CLASS = 'overall'


class Classifier:
    def __init__(self, dv_model_path, lg_model_path, tfidf_model_path=None, mlb_model_path=None):
        self.dv_model = pickle.load(open(dv_model_path, 'rb'))
        self.lg_model = pickle.load(open(lg_model_path, 'rb'))
        self.tfidf_model = pickle.load(open(tfidf_model_path, 'rb'))

    def classify(self, data):
        d = None
        if type(self.dv_model) == CountVectorizer:
            d = self.dv_model.transform([' '.join(data)])
        else:
            d = self.dv_model.transform(get_dict_data(data))

        if self.tfidf_model:
            d = self.tfidf_model.transform(d)
        if self.lg_model is not None:
            if type(self.lg_model) == LinearSVC:
                res = zip(self.lg_model.classes_, self.lg_model._predict_proba_lr(d)[0])
            else:
                res = zip(self.lg_model.classes_, self.lg_model.predict_proba(d)[0])

        # kv = {}
        # for o in res:
        #     k = class_old_class_new_dict.get(o[0])
        #     if k:
        #         kv[k] = kv.get(k, 0.0) + o[1]
        out = sorted(res, key=operator.itemgetter(1), reverse=True)

        return {"category": [o[0] for o in out], "category_prob": [o[1] for o in out]}


class MultiLabelClassifier:
    def __init__(self, dv_model_path, lg_model_path, tfidf_model_path=None, mlb_model_path=None):
        self.dv_model = pickle.load(open(dv_model_path, 'rb'))
        self.lg_model = pickle.load(open(lg_model_path, 'rb'))
        self.tfidf_model = pickle.load(open(tfidf_model_path, 'rb'))
        self.mlb_model = pickle.load(open(mlb_model_path, 'rb'))

    def classify(self, data):
        d = None
        if type(self.dv_model) == CountVectorizer:
            d = self.dv_model.transform([' '.join(data)])
        else:
            d = self.dv_model.transform(get_dict_data(data))

        if self.tfidf_model:
            d = self.tfidf_model.transform(d)
        if self.lg_model is not None:
            labels_d = self.lg_model.predict(d)
            labels = self.mlb_model.inverse_transform(np.array(labels_d))
        return {"category": [labels_d, labels]}


def get_dict_data(data):
    kv = {}
    for d in data:
        kv[d] = kv.get(d, 0.0) + 1.0
    return kv
